// --------------------------------------------------------------
//      Draw Plots for CEvNS Calibration
//
// --------------------------------------------------------------

#include <vector>
#include <string>
#include <regex>
#include <sys/stat.h>
#include <iostream>
#include <variant>

#include "TROOT.h"
#include "TSystem.h"
#include "TString.h"
#include "TRandom.h"
#include "TH1F.h"
#include "TGraph.h"
#include "TChain.h"
#include "TTree.h"
#include "TFile.h"
#include "TBranch.h"
#include "TLeaf.h"
#include "TCanvas.h"
#include "Rtypes.h"

#include "yaml-cpp/yaml.h"

#include "lhcbStyle.h"
#include "OneDimPlot.h"
#include "TwoDimPlot.h"
#include "ScatterPlot.h"

typedef void (*DrawFunction)(std::string, YAML::Node, TChain*);

bool IfFileExist(std::string filePath)
{
    struct stat buffer;   
    return (stat(filePath.c_str(), &buffer) == 0);
}

std::vector<std::string> SplitStr(const char* str, const char* split)
{
    std::string vStr(str);
    std::vector<std::string> vector;
    std::regex delimiter(split);
    std::sregex_token_iterator iter(vStr.begin(), vStr.end(), delimiter, -1);
    std::sregex_token_iterator end;
    while (iter != end)
    {   
        vector.push_back(std::string(*iter));
        ++iter;
    }
    return vector;
}

void DrawPlots(YAML::Node &config, std::vector<std::string> fileList)
{
    //std::string SavePath = config["SavePath"].as<std::string>(".");
    TChain* chain = new TChain(config["TreeName"].as<std::string>("ch0").c_str());
    for(auto filePath:fileList)
    {
        chain->Add(filePath.c_str());
    }

    std::map<std::string, DrawFunction> supportPlot = 
    {
        {"OneDim", &DrawOneDimPlot},
        {"TwoDim", &DrawTwoDimPlot},
        {"Scatter", &DrawScatterPlot},
    };

    for(auto type:supportPlot)
    {
        if(config[type.first].IsDefined())
        {
            for (auto aPlot : config[type.first])
            {
                (*type.second)(aPlot.first.as<std::string>(), aPlot.second, chain);
            }
        }
    }

    delete chain;
}

//main function
int main(int argc, char **argv)
{
    std::string cardPath = std::string(PROJECT_SOURCE_DIR) + std::string("/DrawOptions.yml");
    //std::string cardPath = std::format("{}/Draw.yml", PROJECT_SOURCE_DIR);
    std::vector<std::string> fileList; fileList.clear();
    // Add arglist control
    if (argc >= 2)
    {
        for(int iA = 1; iA < argc; ++iA)
        {
            if(SplitStr(argv[iA], "\\.").back() == std::string("yml"))
            {
                cardPath = TString(argv[iA]);
            }
            if(SplitStr(argv[iA], "\\.").back() == std::string("root"))
            {
                std::cout << "Find root file: " << std::string(argv[iA]) << std::endl;
                if (!IfFileExist(std::string(argv[iA])))
                {
                    std::clog << "Warning: ROOT file \" " << cardPath << " \" NOT FOUND." << std::endl;
                    continue;
                }
                fileList.push_back(std::string(argv[iA]));
            }
        }
    }
    std::cout << "Using YAML config: " << cardPath << std::endl;
    if (!IfFileExist(cardPath))
    {
        std::cerr << "Error: YAML file \" " << cardPath << " \" NOT FOUND!!! EXIT." << std::endl;
        return EXIT_FAILURE;
    }
    if(fileList.size() == 0)
    {
        std::cerr << "Error: No valid ROOT file!!! EXIT." << std::endl;
        return EXIT_FAILURE;
    }

    //Load LHCb style
    lhcbStyle();

    //Load yaml config
    YAML::Node config = YAML::LoadFile(cardPath);

    //Draw plots
    DrawPlots(config, fileList);

    return EXIT_SUCCESS;
}
